import cv2
import numpy as np

class CLAHEPreprocessor:
    def __init__(self, clip_limit=2.0, tile_grid_size=(8, 8), brightness_balance=None):
        """
        初始化 CLAHE 预处理器
        :param clip_limit: CLAHE算法中的限制对比度因子，默认为2.0
        :param tile_grid_size: CLAHE算法中的网格大小，默认为8x8
        :param brightness_balance: 亮度平衡因子，如果为None，则不进行亮度平衡
        """
        self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
        self.brightness_balance = brightness_balance

    # def __call__(self, left_image, right_image):
    #     """
    #     读取左右眼图像，并应用 CLAHE 进行对比度增强、亮度平衡和去噪
    #     :param left_image: 左眼图像
    #     :param right_image: 右眼图像
    #     :return: 返回增强后的左右眼图像
    #     """
    #     # 对左眼图像进行处理
    #     left_img = self._process_image(left_image)
    #     # 对右眼图像进行处理
    #     right_img = self._process_image(right_image)
    #
    #     return left_img, right_img
    def __call__(self, image):
        """
        读取图像，并应用 CLAHE 进行对比度增强、亮度平衡和去噪
        :return: 返回增强后的图像
        """
        # 对左眼图像进行处理
        processed_image = self._process_image(image)
        return processed_image

    def _process_image(self, image):
        """
        对单张图像进行处理，包括CLAHE对比度增强、亮度平衡和去噪
        :param image: 输入图像
        :return: 处理后的图像
        """
        # 如果设置了亮度平衡因子，则调整图像亮度
        if self.brightness_balance:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            brightness = gray.sum() / gray.size
            image = np.uint8(np.minimum(image * self.brightness_balance / brightness, 255))

        # 应用 CLAHE 对比度增强
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        lab_planes = list(cv2.split(lab))  # 将元组转换为列表
        lab_planes[0] = self.clahe.apply(lab_planes[0])  # 对亮度通道进行CLAHE处理
        lab = cv2.merge(lab_planes)  # 重新组合为图像
        image = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

        # 应用去噪
        image = cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 1, 3)
        image = cv2.bilateralFilter(image, 5, 1, 1)

        return image


if __name__ == '__main__':
    img = cv2.imread('示例.jpg')
    clahe = CLAHEPreprocessor()
    processed_img = clahe(img)
    cv2.imshow('img', processed_img)
